from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Random import get_random_bytes
from Crypto.Hash import SHA512
import base64
import json
from typing import Dict, Union, Optional, Any


class EnhancedAES256Encryptor:
    """增强的AES-256加密/解密工具，支持字节到字节的加密解密和旧版字典格式"""
    
    # 加密数据格式版本
    CURRENT_VERSION = '2.0'
    BINARY_VERSION = '3.0'  # 二进制格式版本

    def __init__(self, master_key: bytes):
        self.iterations = 10000
        self._master_key = master_key

    def _derive_key(self, salt: bytes) -> bytes:
        """密钥派生函数（包含盐值处理）"""
        return PBKDF2(
            self._master_key, 
            salt, 
            dkLen=32, 
            count=self.iterations,
            hmac_hash_module=SHA512
        )

    def encrypt(self, plaintext_bytes: bytes, binary_output: bool = True) -> Union[bytes, Dict[str, Any]]:
        """
        加密流程（包含盐值生成）
        
        参数:
            plaintext_bytes: 要加密的明文字节
            binary_output: 是否输出二进制格式（用于文件加密）
            
        返回:
            如果binary_output=True，返回加密后的字节
            否则返回包含加密数据的字典（向后兼容）
        """
        salt = get_random_bytes(16)
        key = self._derive_key(salt)
        cipher = AES.new(key, AES.MODE_GCM)
        ciphertext, tag = cipher.encrypt_and_digest(plaintext_bytes)
        
        # 创建包含加密数据的字典
        encrypted_data = {
            'version': self.BINARY_VERSION if binary_output else self.CURRENT_VERSION,  # 版本号
            'ciphertext': ciphertext,  # 加密后的数据
            'salt': salt,              # 盐值
            'nonce': cipher.nonce,     # 随机数
            'tag': tag,                # 认证标签
            'iterations': self.iterations  # 迭代次数
        }
        
        if binary_output:
            # 创建二进制格式：
            # 1. 将非binary字段base64编码为字符串
            binary_data = {
                'version': self.BINARY_VERSION,
                'ciphertext': base64.b64encode(ciphertext).decode(),
                'salt': base64.b64encode(salt).decode(),
                'nonce': base64.b64encode(cipher.nonce).decode(),
                'tag': base64.b64encode(tag).decode(),
                'iterations': self.iterations
            }
            # 2. 将整个字典序列化为JSON字符串
            json_str = json.dumps(binary_data)
            # 3. 将JSON字符串编码为字节
            return json_str.encode('utf-8')
        else:
            # 向后兼容：返回base64编码的字典
            return {
                'version': self.CURRENT_VERSION,  
                'ciphertext': base64.b64encode(ciphertext).decode(),
                'salt': base64.b64encode(salt).decode(),
                'nonce': base64.b64encode(cipher.nonce).decode(),
                'tag': base64.b64encode(tag).decode(),
                'iterations': self.iterations
            }

    def decrypt(self, encrypted_data: Union[bytes, Dict[str, Any]]) -> bytes:
        """
        解密流程
        
        参数:
            encrypted_data: 加密数据，可以是二进制格式或字典
            
        返回:
            解密后的字节数据
        """
        # 处理不同类型的输入
        if isinstance(encrypted_data, bytes):
            try:
                # 尝试将字节解析为JSON
                binary_data = json.loads(encrypted_data.decode('utf-8'))
                # 解码base64字段
                salt = base64.b64decode(binary_data['salt'])
                nonce = base64.b64decode(binary_data['nonce'])
                tag = base64.b64decode(binary_data['tag'])
                ciphertext = base64.b64decode(binary_data['ciphertext'])
            except Exception as e:
                raise ValueError(f"无法解析二进制加密数据: {e}")
        elif isinstance(encrypted_data, dict):
            # 参数校验
            required_fields = {'ciphertext', 'salt', 'nonce', 'tag'}
            if not all(field in encrypted_data for field in required_fields):
                raise ValueError("无效的加密数据格式")

            # 解码参数
            salt = base64.b64decode(encrypted_data['salt'])
            nonce = base64.b64decode(encrypted_data['nonce'])
            tag = base64.b64decode(encrypted_data['tag'])
            ciphertext = base64.b64decode(encrypted_data['ciphertext'])
        else:
            raise TypeError(f"不支持的加密数据类型: {type(encrypted_data)}")

        # 密钥派生
        key = self._derive_key(salt)
        
        # 解密验证
        cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
        return cipher.decrypt_and_verify(ciphertext, tag)
        
    def decrypt_to_string(self, encrypted_data: Union[bytes, Dict[str, Any]]) -> str:
        """
        解密并返回字符串（向后兼容方法）
        
        参数:
            encrypted_data: 加密数据，可以是二进制格式或字典
            
        返回:
            解密后的UTF-8字符串
        """
        decrypted_bytes = self.decrypt(encrypted_data)
        return decrypted_bytes.decode('utf-8')
